/*
 * Copyright (c) 2016, Roman Meyta <theshrodingerscat@gmail.com>
 * Copyright (c) 2020-2021 https://gitee.com/fsfzp888
 * All rights reserved
 */

#include <QByteArray>
#include <QMessageBox>
#include <QObject>

#include <fstream>

#include "ImageFormats.h"
#include "Logger.h"
#include "jpeglib.h"
#include "turbojpeg.h"

/**
 * @brief Make QImage in specified format
 *
 * @param data
 * @param len
 * @param width
 * @param height
 * @return QImage
 */
QImage simpleQimage(const unsigned char *data, unsigned, unsigned width, unsigned height, QImage::Format format)
{
    return QImage(data, width, height, format);
}

QImage makeARGB32Image(const unsigned char *data, unsigned len, unsigned width, unsigned height)
{
    return simpleQimage(data, len, width, height, QImage::Format_ARGB32).rgbSwapped();
}

QImage makeRGB32Image(const unsigned char *data, unsigned len, unsigned width, unsigned height)
{
    return simpleQimage(data, len, width, height, QImage::Format_RGB32).rgbSwapped();
}

QImage makeRGB24Image(const unsigned char *data, unsigned len, unsigned width, unsigned height)
{
    return simpleQimage(data, len, width, height, QImage::Format_RGB888).rgbSwapped();
}

QImage makeRGB555Image(const unsigned char *data, unsigned len, unsigned width, unsigned height)
{
    return simpleQimage(data, len, width, height, QImage::Format_RGB555).rgbSwapped();
}

QImage makeRGB565Image(const unsigned char *data, unsigned len, unsigned width, unsigned height)
{
    return simpleQimage(data, len, width, height, QImage::Format_RGB16).rgbSwapped();
}

QImage makeJPGImage(const unsigned char *data, unsigned len, unsigned, unsigned)
{
    QByteArray bytes(reinterpret_cast<const char *>(data), len);
    QImage image;
    image.loadFromData(bytes, "JPEG");
    return image;
}

static int YUV2RGB(const unsigned char *pYUV, unsigned char *pRGB, int width, int height)
{
    if (nullptr == pYUV || nullptr == pRGB)
    {
        return -1;
    }
    const unsigned char *pYUVData = pYUV;
    unsigned char *pRGBData       = pRGB;
    int Y1, U1, V1, Y2, R1, G1, B1, R2, G2, B2;
    int C1, D1, E1, C2;
    for (int i = 0; i < height; ++i)
    {
        for (int j = 0; j < width / 2; ++j)
        {
            Y1 = *(pYUVData + i * width * 2 + j * 4);
            U1 = *(pYUVData + i * width * 2 + j * 4 + 1);
            Y2 = *(pYUVData + i * width * 2 + j * 4 + 2);
            V1 = *(pYUVData + i * width * 2 + j * 4 + 3);
            C1 = Y1 - 16;
            C2 = Y2 - 16;
            D1 = U1 - 128;
            E1 = V1 - 128;
            R1 = ((298 * C1 + 409 * E1 + 128) >> 8 > 255 ? 255 : (298 * C1 + 409 * E1 + 128) >> 8);
            G1 = ((298 * C1 - 100 * D1 - 208 * E1 + 128) >> 8 > 255 ? 255 : (298 * C1 - 100 * D1 - 208 * E1 + 128) >> 8);
            B1 = ((298 * C1 + 516 * D1 + 128) >> 8 > 255 ? 255 : (298 * C1 + 516 * D1 + 128) >> 8);
            R2 = ((298 * C2 + 409 * E1 + 128) >> 8 > 255 ? 255 : (298 * C2 + 409 * E1 + 128) >> 8);
            G2 = ((298 * C2 - 100 * D1 - 208 * E1 + 128) >> 8 > 255 ? 255 : (298 * C2 - 100 * D1 - 208 * E1 + 128) >> 8);
            B2 = ((298 * C2 + 516 * D1 + 128) >> 8 > 255 ? 255 : (298 * C2 + 516 * D1 + 128) >> 8);
            *(pRGBData + i * width * 3 + j * 6 + 2) = R1 < 0 ? 0 : R1;
            *(pRGBData + i * width * 3 + j * 6 + 1) = G1 < 0 ? 0 : G1;
            *(pRGBData + i * width * 3 + j * 6)     = B1 < 0 ? 0 : B1;
            *(pRGBData + i * width * 3 + j * 6 + 5) = R2 < 0 ? 0 : R2;
            *(pRGBData + i * width * 3 + j * 6 + 4) = G2 < 0 ? 0 : G2;
            *(pRGBData + i * width * 3 + j * 6 + 3) = B2 < 0 ? 0 : B2;
        }
    }
    return 0;
}

unsigned long YUV2ToJPG(int width, int height, const unsigned char *inputYuv, unsigned char **outJpeg)
{
    struct jpeg_compress_struct cinfo;
    struct jpeg_error_mgr jerr;
    JSAMPROW row_pointer[1];
    int i = 0, j = 0;
    const unsigned char *pY, *pU, *pV;
    unsigned char *yuvbuf = new unsigned char[width * 3];
    cinfo.err             = jpeg_std_error(&jerr);
    jpeg_create_compress(&cinfo);
    unsigned long outSize = 0;
    jpeg_mem_dest(&cinfo, outJpeg, &outSize);
    cinfo.image_width      = width;
    cinfo.image_height     = height;
    cinfo.input_components = 3;
    cinfo.in_color_space   = JCS_YCbCr;
    cinfo.dct_method       = JDCT_FLOAT;
    jpeg_set_defaults(&cinfo);
    jpeg_set_quality(&cinfo, 100, TRUE);
    jpeg_start_compress(&cinfo, TRUE);
    pY = inputYuv;
    pU = inputYuv + 1;
    pV = inputYuv + 3;
    j  = 1;
    while (cinfo.next_scanline < cinfo.image_height)
    {
        int index = 0;
        for (i = 0; i < width; i += 2)
        {
            yuvbuf[index++] = *pY;
            yuvbuf[index++] = *pU;
            yuvbuf[index++] = *pV;
            pY += 2;
            yuvbuf[index++] = *pY;
            yuvbuf[index++] = *pU;
            yuvbuf[index++] = *pV;
            pY += 2;
            pU += 4;
            pV += 4;
        }
        row_pointer[0] = yuvbuf;
        jpeg_write_scanlines(&cinfo, row_pointer, 1);
        j++;
    }
    jpeg_finish_compress(&cinfo);
    jpeg_destroy_compress(&cinfo);
    delete[] yuvbuf;
    return outSize;
}

QImage makeYUY2Image(const unsigned char *data, unsigned, unsigned width, unsigned height)
{
    // unsigned char *jpeg_buf = nullptr;
    // unsigned long jpeg_size = YUV2ToJPG(width, height, data, &jpeg_buf);
    // if (jpeg_size)
    //{
    //    std::ofstream ofs("tmp.jpg", std::ios_base::out|std::ios_base::trunc|std::ios_base::binary);
    //    ofs.write((const char*)jpeg_buf, jpeg_size);
    //    ofs.close();
    //    QByteArray bytes(reinterpret_cast<const char*>(jpeg_buf), jpeg_size);
    //    QImage image;
    //    image.loadFromData(bytes, "JPEG");
    //    delete[] jpeg_buf;
    //    return image;
    //}
    unsigned int rgb_len   = width * height * 3;
    unsigned char *rgb_buf = new unsigned char[rgb_len];
    YUV2RGB(data, rgb_buf, width, height);
    QImage img = simpleQimage(rgb_buf, rgb_len, width, height, QImage::Format_RGB888).rgbSwapped();
    delete[] rgb_buf;
    return img;
}

QImage dummy(const unsigned char *, unsigned, unsigned, unsigned)
{
    QMessageBox mb;
    mb.setText(QObject::tr("current image format not supported"));
    mb.exec();
    QImage image;
    return image;
}

const std::string &getImageFormatName(const GUID &uid)
{
    static std::string defaultName = "unknown";
    for (auto &formatRow : ImageFormatTable)
    {
        if (formatRow.directshowFormat == uid)
        {
            return formatRow.name;
        }
    }
    return defaultName;
}

QImageMaker getQImageMaker(const GUID &uid)
{
    for (auto &formatRow : ImageFormatTable)
    {
        if (formatRow.directshowFormat == uid)
        {
            return formatRow.makeQImage;
        }
    }
    LOG_ERROR("Image format %s not supported", getImageFormatName(uid));
    return dummy;
}

bool isKnownImageFormat(const GUID &uid)
{
    for (auto &formatRow : ImageFormatTable)
    {
        if (formatRow.directshowFormat == uid)
        {
            return true;
        }
    }
    return false;
}
